#pragma once

#include "DBN.h"

//#include "../Loirey/loirey_FFT.h"

using namespace loirey;

namespace nw_rbm
{
	class CConvolutionalDBN_for_Image : public CDeepBeliefNetwork
	{
	public:
		class CBasicLayerInfo
		{
		public:
			int K_Visible, K_Hidden;		// Feature Dimension
			int wMask, hMask, aMask, sMask;	// Feature Window Size
			int xPool, yPool;				// Max Pooling Window Size

		public:
			virtual void myInit(
				int K_Visible, int K_Hidden,
				int wMask, int hMask,
				int xPool, int yPool);
			virtual void LoadFromStream(istream& inStream);
			virtual void LoadFromStream(istream& inStream, int K_Visible);
			virtual void SaveToStream(ostream& outStream) const;
		};

	public:
		class CExtraLayerInfo : public CBasicLayerInfo
		{
		public:
			int wVisible, hVisible, aVisible, sVisible;
			int wHidden, hHidden, aHidden, sHidden;
			int wPool, hPool, aPool, sPool;

		public:
			CExtraLayerInfo();
			virtual void Clear();
		public:
			virtual void Init_RBM_InferenceHelper(
				int LayerIndex,
				C_RBM_InferenceHelper*& p_RBM_InferenceHelper);
		public:
			virtual int GetAmountConnection_Visible(
				int UnitHandle) const;
			virtual int GetUnitConnectionList_Visible(
				int UnitHandle,
				CSimpleTypeArray<int>& vConnectedUnitHandle,
				CSimpleTypeArray<int>& vEdgeHandle) const;
			virtual int GetAmountConnection_Hidden(
				int UnitHandle) const;
			virtual int GetUnitConnectionList_Hidden(
				int UnitHandle,
				CSimpleTypeArray<int>& vConnectedUnitHandle,
				CSimpleTypeArray<int>& vEdgeHandle) const;
		};

	public:
		class C_RBM_InferenceHelper_Convolution : public C_RBM_InferenceHelper
		{
		public:
			const CExtraLayerInfo* qExtraLayerInfo;

		public:
			virtual ~C_RBM_InferenceHelper_Convolution();
			virtual void InferenceEnergy_of_Units_Visible(CRestrictedBoltzmannMachine* pRBM);
			virtual void InferenceEnergy_of_Units_Hidden(CRestrictedBoltzmannMachine* pRBM);
			virtual RBM_FLOAT ComputeEdgeGradient_ConDiv(
				CRestrictedBoltzmannMachine* pRBM,
				const RBM_FLOAT* qVisibleState_Data,
				const RBM_FLOAT* qHiddenExpectation_Data,
				const RBM_FLOAT* qVisibleState_Model,
				const RBM_FLOAT* qHiddenExpectation_Model,
				RBM_FLOAT ExampleWeight, RBM_FLOAT GradientAmendment);
		};

	public:
		class CSettings_Layer : public CBasicLayerInfo
		{
		public:
		// For RBM Learning
			RBM_FLOAT SpaCod_SparseLevel;		
			RBM_FLOAT SpaCod_Lambda;
			RBM_FLOAT ConDiv_Decay_Mask;
			RBM_FLOAT ConDiv_Decay_Bias;
			RBM_FLOAT ConDiv_Momentum;
			RBM_FLOAT LearningRate;
		public:
		// Layer Details
			int Input_xPool;
			int Input_yPool;
			int Output_xPool;
			int Output_yPool;
			int wFullMask_Visible;
			int hFullMask_Visible;
			int wFullMask_Hidden;
			int hFullMask_Hidden;
			int wFullMask_Pooling;
			int hFullMask_Pooling;

		public:
			virtual void SetInitialLayerDetails();
			virtual void SetRegularLayerDetails(const CSettings_Layer& PreviousLayer);
		public:
			virtual void LoadFromStream(istream& inStream);
			virtual void LoadFromStream(istream& inStream, int K_Visible);
			virtual void SaveToStream(ostream& outStream) const;
		public:
			virtual void InitWeightItemSettings_for_RBM(
				CRestrictedBoltzmannMachine* pRBM);
			virtual void InitUnitSettings_for_RBM(
				CRestrictedBoltzmannMachine* pRBM, CExtraLayerInfo& ExtraLayerInfo,
				int wSource, int hSource, RBM_FLOAT GaussianSigma_Visible);
			virtual void InitUnitSettings_for_MP_RBM(
				CRBM_with_MaxPooling* pMP_RBM, CExtraLayerInfo& ExtraLayerInfo,
				int wSource, int hSource, RBM_FLOAT GaussianSigma_Visible);
		};

	public:
		class CSettings_CDBN
		{
		public:
			int Depth;
			int K_Visible;
			RBM_FLOAT GaussianSigma_Visible;
			CSimpleTypeArray<CSettings_Layer> ListLayerSettings;

		public:
			CSettings_CDBN();
			virtual void Clear();
			virtual void LoadFromFile(string strFN_CDBN_Settings);
			//virtual void SaveToFile(string strFN_CDBN_Settings) const;
		};

	public:
		class CLayerMaskImages
		{
		public:
			int W, H, V, U, A, B, SS;
			CSimpleTypeArray<RBM_FLOAT> vmatData;
			CSimpleTypeArray<RBM_FLOAT> vmatMean;
			CSimpleTypeArray<RBM_FLOAT> vmatVar;
		};
	public:

	public:
		CSettings_CDBN Settings;
		CSimpleTypeArray<CExtraLayerInfo> ListExtraLayerInfo;
		CSimpleTypeArray<CLayerMaskImages> vLayerMaskImage_Hidden;
		CSimpleTypeArray<CLayerMaskImages> vLayerMaskImage_Pooling;

	public:
		CConvolutionalDBN_for_Image();
		virtual void Clear();
		virtual void DeepClear();
		CExtraLayerInfo& GetExtraLayerInfo(int LayerIndex);
		CSettings_Layer& GetSingleLayerSettings(int LayerIndex);
		CLayerMaskImages& GetSingleLayerMaskImages_Hidden(int LayerIndex);
		CLayerMaskImages& GetSingleLayerMaskImages_Pooling(int LayerIndex);
	public:
		virtual void RefreshLayerConnectionManagers(int LayerIndex);
	public:
		virtual void Make_RBM_InferenceHelper(
			int LayerIndex,
			C_RBM_InferenceHelper*& p_RBM_InferenceHelper);
	public:
		void GetLayerCumulativePoolingSize_VisibleUnits(int LayerIndex, int& xPool, int& yPool) const;
		void GetLayerCumulativePoolingSize_PoolingUnits(int LayerIndex, int& xPool, int& yPool) const;
	public:
		virtual void InitBySettings();
		virtual void InitByInputSize(int W, int H);
	public:
		virtual void SetVisibleState(int DimensionIndex, const RBM_FLOAT* qSrc);
		virtual const CUnitBasic* GetUnits_Visible(
			int LayerIndex, int VisibleDimensionIndex,
			int& DstW, int& DstH);
		virtual const CUnitBasic* GetUnits_Hidden(
			int LayerIndex, int HiddenDimensionIndex,
			int& DstW, int& DstH);
		virtual const CUnitPooling* GetUnits_Pooling(
			int LayerIndex, int HiddenDimensionIndex,
			int& DstW, int& DstH);
		virtual const CWeightItem* GetWeightMatrix(
			int LayerIndex, int VisibleDimensionIndex, int HiddenDimensionIndex);
	//public:
	//	virtual RBM_FLOAT ComputeWeightMatrixValueRange(
	//		int LayerIndex);
	//	virtual void DrawWeightMatrix(
	//		int LayerIndex, int VisibleDimensionIndex, int HiddenDimensionIndex,
	//		RBM_FLOAT ZoomScale, CImg<RBM_FLOAT>& imgDst);
	//	virtual void DrawNormalizedWeightMatrix(
	//		int LayerIndex, int VisibleDimensionIndex, int HiddenDimensionIndex,
	//		RBM_FLOAT ZoomScale, RBM_FLOAT MaxValue,	CImg<int>& imgDst);
	public:
		virtual void ComputeMaskImages(
			int LayerIndex);
		virtual void DrawMaskImage(
			int LayerIndex, int ChannelIndex, int GroupIndex,
			RBM_FLOAT ZoomScale, CImg<RBM_FLOAT>& imgDst);
		virtual void DrawMaskImage(
			int LayerIndex, int ChannelIndex, int GroupIndex,
			RBM_FLOAT ZoomScale, RBM_FLOAT MinValue, RBM_FLOAT MaxValue, CImg<int>& imgDst);
		virtual void DrawMaskImage(
			int LayerIndex, int ChannelIndex, int GroupIndex,
			RBM_FLOAT ZoomScale, CImg<int>& imgDst);
	public:
		static const int DRAW_UNITS_PROPERTY_MODE_VisibleState = 0;
		static const int DRAW_UNITS_PROPERTY_MODE_HiddenExpectation = 4;
		static const int DRAW_UNITS_PROPERTY_MODE_PoolingExpectation = 5;
		virtual void DrawUnitsPropertyImage(
			int LayerIndex, int Mode, int DimensionIndex,
			RBM_FLOAT ZoomScale, CImg<RBM_FLOAT>& imgDst);
		virtual void DrawUnitsPropertyImage(
			int LayerIndex, int Mode, int DimensionIndex,
			RBM_FLOAT ZoomScale, CImg<int>& imgDst);
		virtual void DrawAllUnitsPropertyImage(
			int LayerIndex, int Mode,
			RBM_FLOAT ZoomScale, int xDraw, CImg<int>& imgDst);
	public:
		virtual void ComputeBothGradient(
			int LayerIndex,
			RBM_FLOAT SpaCod_SparseLevel,
			RBM_FLOAT ExampleWeight, RBM_FLOAT GradientAmendment);
		virtual void UpdateBothGradient(
			int LayerIndex,
			RBM_FLOAT LearningRate, RBM_FLOAT SpaCod_Lambda,
			RBM_FLOAT ConDiv_Decay_Mask, RBM_FLOAT ConDiv_Decay_Bias, RBM_FLOAT ConDiv_Momentum);
		virtual RBM_FLOAT GetCumulativeCosts(
			int LayerIndex,
			RBM_FLOAT SpaCod_Lambda,
			RBM_FLOAT& VisibleError, RBM_FLOAT& HiddenPenalty);
	public:
		virtual void ComputeBothGradient(
			int LayerIndex,
			RBM_FLOAT ExampleWeight);
		virtual void UpdateBothGradient(
			int LayerIndex);
		virtual RBM_FLOAT GetCumulativeCosts(
			int LayerIndex,
			RBM_FLOAT& VisibleError, RBM_FLOAT& HiddenPenalty);
	};
}

